import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA


register_matplotlib_converters()


class Beili(object):
    def __init__(self, codes, start=None, end=None, valid_bars=5, macd_limit=0.05,
                 market_type=QA.MARKET_TYPE.STOCK_CN, frequence=QA.FREQUENCE.DAY):
        if isinstance(codes, list):
            self.codes = codes
            self.names = None
        else:
            self.codes = codes['code'].to_list()
            self.names = codes[['code', 'name']].set_index(['code'])
        data = QA.QA_quotation(self.codes, start, end, frequence=frequence, market=market_type,
                                        source=QA.DATASOURCE.MONGO, output=QA.OUTPUT_FORMAT.DATASTRUCT)
        if market_type == QA.MARKET_TYPE.STOCK_CN:
            data = data.to_qfq()
        self.data = data.data
        self.valid_bars = valid_bars
        self.macd_limit = macd_limit

    def radar(self):
        for code in self.codes:
            df = self.data.xs(code, level='code', axis=0, drop_level=True)
            df = QA.QA_indicator_beili(df, valid_bars=self.valid_bars, macd_limit=self.macd_limit).fillna(0)
            ind_last_last_bar = df.iloc[-3]
            ind_last_bar = df.iloc[-2]
            ind_current_bar = df.iloc[-1]
            # hbl = ind_current_bar.H_BL + ind_last_bar.H_BL + ind_last_last_bar.H_BL
            # lbl = ind_current_bar.L_BL + ind_last_bar.L_BL + ind_last_last_bar.L_BL
            hbl = ind_current_bar.H_BL
            lbl = ind_current_bar.L_BL
            if hbl > 0:
                print('code：{}---{}---顶背离---时间 {}---CLOSE {}---LEVEL {}'.format(code, self.names.loc[code]['name'],
                                                                                ind_current_bar.name,
                                                                                ind_current_bar.CLOSE, hbl))
                #if hbl == 3:
                    #self._plot(df)
                # self._plot(df)
            if lbl > 0:
                print('code：{}---{}---底背离---时间 {}---CLOSE {}---LEVEL {}'.format(code, self.names.loc[code]['name'],
                                                                                ind_current_bar.name,
                                                                                ind_current_bar.CLOSE, lbl))
                #if lbl == 3:
                    #self._plot(df)
                # self._plot(df)

    def plot_single(self, code='000001'):
        df = self.data.xs(code, level='code', axis=0, drop_level=True)
        data = QA.QA_indicator_beili(df, valid_bars=self.valid_bars, macd_limit=self.macd_limit).fillna(0)
        self._plot(data)

    def _plot(self, data):
        figure = plt.figure()
        axes1 = figure.add_subplot(411)
        axes2 = figure.add_subplot(412)
        axes3 = figure.add_subplot(413)
        axes4 = figure.add_subplot(414)
        axes1.plot(data.index, data['CLOSE'])
        axes2.plot(data.index, data[['DIFF', 'DEA']])
        axes3.bar(data.index, data['MACD'])
        axes4.plot(data.index, data[['H_BL', 'L_BL']])

        left, right = axes3.get_xlim()
        axes3.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
        multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
        plt.show()


if __name__ == '__main__':
    code = '601933'
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    fetch_data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    fetch_data = fetch_data[fetch_data.symbol.str.startswith('00')
                      | fetch_data.symbol.str.startswith('60')
                      | fetch_data.symbol.str.startswith('30')]
    fetch_data = fetch_data[fetch_data.list_date < '20190101']
    #fetch_data = fetch_data[fetch_data.symbol.str.startswith(code)]
    all = fetch_data.rename(columns={'symbol': 'code'})
    stocks = all[~all.name.str.contains('ST')]

    stocks = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\EXAMPLE\\AI\\data\\hs300.csv', dtype=str)
    # hs300 = ts.get_hs300s()

    # index = [code]
    # bl = Beili(index, start='2020-08-01', end='2021-01-29', frequence=QA.FREQUENCE.SIXTY_MIN, market_type=QA.MARKET_TYPE.INDEX_CN)

    bl = Beili(stocks, start='2020-08-01', end=str(datetime.date.today()), frequence=QA.FREQUENCE.DAY, market_type=QA.MARKET_TYPE.STOCK_CN)
    # bl.plot_single(code)
    bl.radar()
